import torch

from backbones.rnn import RNN
from datasets.dataset import generate_weather_loader
from torch import nn
import torch

if __name__ == '__main__':
    train_loader = generate_weather_loader(seq_len=10, batch_size=20)
    vocab_size = 46  # 决定天气因素的特征数量
    hidden_size = 20
    batch_size = 20
    device = torch.device("cuda")
    model = RNN(vocab_size, hidden_size, 2, device=device)

    criterion_mse = nn.MSELoss()
    criterion_ce = nn.CrossEntropyLoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    epochs = 1000
    for epoch in range(epochs):
        loss_list = []
        for inputs, outputs in train_loader:
            inputs = inputs.to(device)
            outputs = outputs.to(device)
            optimizer.zero_grad()
            h = torch.zeros((2, batch_size, hidden_size), device=device)
            predicts, _ = model(inputs, h)

            # 气温与风速的均方差损失
            loss1 = criterion_mse(predicts[:, :6], outputs[:, :, :6].reshape(-1, 6))

            # 白天天气分类交叉熵损失
            loss2 = criterion_ce(predicts[:, 6:18], outputs.reshape(-1, vocab_size)[:, 6:18])

            # 夜晚天气分类交叉熵损失
            loss3 = criterion_ce(predicts[:, 18:30], outputs.reshape(-1, vocab_size)[:, 18:30])

            # 白天风速分类交叉熵损失
            loss4 = criterion_ce(predicts[:, 30:38], outputs.reshape(-1, vocab_size)[:, 30:38])

            # 夜晚风速分类交叉熵损失
            loss5 = criterion_ce(predicts[:, 38:], outputs.reshape(-1, vocab_size)[:, 38:])

            # 联合损失
            loss = loss1 + loss2 + loss3 + loss4 + loss5
            loss.backward()
            optimizer.step()

            loss_list.append(loss.item())
        avg_loss = sum(loss_list) / len(loss_list)
        print(f"epoch:{epoch + 1}/{epochs} -- loss:{avg_loss:.4f}")
